package com.demo.datastructure.tree;

/**
 * @description: 简单红黑树的实现
 * @author: ljr
 * @create: 2022/1/26 4:03 PM
 */
public class SimpleRedBlackTree<Key extends Comparable<Key>, Value> {

    private Node root;

    private static final boolean RED = Boolean.TRUE;

    private static final boolean BLACK = Boolean.FALSE;

    private boolean isRed(Node node){
        if(node == null) return BLACK;
        return node.color == RED;
    }

    private Node rotateLeft(Node h){
         Node right = h.right;
         h.right = right.left;
         right.left = h;
         right.color = h.color;
         h.color = RED;
         right.n = size(h) + size(right.right) + 1;
         return right;
    }

    private Node rotateRight(Node h){
        Node left = h.left;
        h.left = left.right;
        left.right = h;
        left.color = h.color;
        h.color = RED;
        left.n = size(h) + size(left.left) + 1;
        return left;
    }


    public Node put(Key key, Value value){
        Node newNode = new Node(key, value, 1, RED);
        root = put(root, newNode);
        root.color = BLACK;
        return newNode;
    }

    private Node put(Node parentNode, Node node) {
        if(parentNode == null){
            return node;
        }
        int cmp = parentNode.key.compareTo(node.key);
        if(cmp > 0){
            parentNode.left = put(parentNode.left, node);
        }else if(cmp < 0){
            parentNode.right = put(parentNode.right, node);
        }else{
            parentNode.value = node.value;
        }

        if(isRed(parentNode.right) && !isRed(parentNode.left)) parentNode = rotateLeft(parentNode);
        if(isRed(parentNode.left) && isRed(parentNode.left.left)) parentNode = rotateRight(parentNode);
        if(isRed(parentNode.left) && isRed(parentNode.right)) fliColors(parentNode);

        parentNode.n = size(parentNode.left) + size(parentNode.right) + 1;
        return parentNode;
    }

    /**
     * 上色
     */
    private void fliColors(Node h){
        h.color = RED;
        h.left.color = BLACK;
        h.right.color = BLACK;
    }

    public int level(){
        return level(root);
    }

    private int level(Node node){
        if(node == null) return 0;
        return Math.max(level(node.left), level(node.right)) + 1;
    }

    public int size(){
        return size(root);
    }

    private int size(Node node){
        if(node == null) return 0;
        return node.n;
    }

    private static class Node<Key extends Comparable<Key>, Value>{
        private Key key;

        private Value value;

        private int n;

        private Node left;

        private Node right;

        private boolean color;

        public Node(Key key, Value value, int n, boolean color) {
            this.key = key;
            this.value = value;
            this.n = n;
            this.color = color;
        }

        @Override
        public String toString() {
            return "Node{" +
                    "key=" + key +
                    ", value=" + value +
                    ", n=" + n +
                    ", color=" + color +
                    ", left=" + left +
                    ", right=" + right +
                    '}';
        }
    }


    public static void main(String[] args) {
        SimpleRedBlackTree<String, String> myBinarySearchTree1 = new SimpleRedBlackTree();
        myBinarySearchTree1.put("S", "S");
        myBinarySearchTree1.put("E", "E");
        myBinarySearchTree1.put("A", "A");
        myBinarySearchTree1.put("R", "R");
        myBinarySearchTree1.put("C", "C");
        myBinarySearchTree1.put("H", "H");
        myBinarySearchTree1.put("X", "X");
        myBinarySearchTree1.put("M", "M");
        myBinarySearchTree1.put("P", "P");
        myBinarySearchTree1.put("L", "L");
        System.out.println(myBinarySearchTree1.root);
        System.out.println(myBinarySearchTree1.level());
    }

}
